import sys
import cv2

import numpy as np

import common.img_utils.fh.GraphOperator as go

from common.img_utils.fh.fh_label_map_generator import (
    generate_label_image,
    weight_from_edge,
)


def main():
    sigma = float(sys.argv[1])
    k = float(sys.argv[2])
    min_size = float(sys.argv[3])

    img = cv2.imread(sys.argv[4])
    float_img = np.asarray(img, dtype=float)

    gaussian_img = cv2.GaussianBlur(float_img, (5, 5), sigma)
    b, g, r = cv2.split(gaussian_img)
    smooth_img = (r, g, b)

    height, width, channel = img.shape
    graph = go.build_graph(smooth_img, width, height)

    sorted_graph = sorted(graph, key=weight_from_edge)

    ufset = go.segment_graph(sorted_graph, width * height, k)
    ufset = go.remove_small_component(ufset, sorted_graph, min_size)

    save_img = generate_label_image(ufset, width, height)
    cv2.imwrite(sys.argv[5], save_img)


if __name__ == "__main__":
    main()
